#include "gemm.cuh"
#include "common.cuh"

namespace runtime {namespace cuda {

template <>
void caffe_gpu_gemm<float>(const cublasOperation_t TransA,
    const cublasOperation_t TransB, const int M, const int N, const int K,
    const float alpha, const float* A, const float* B, const float beta,
    float* C, cublasHandle_t& cublas_handle_) {
  // Note that cublas follows fortran order.
  int lda = (TransA == CUBLAS_OP_N) ? K : M;
  int ldb = (TransB == CUBLAS_OP_N) ? N : K;
  cublasOperation_t cuTransA =
      (TransA == CUBLAS_OP_N) ? CUBLAS_OP_N : CUBLAS_OP_T;
  cublasOperation_t cuTransB =
      (TransB == CUBLAS_OP_N) ? CUBLAS_OP_N : CUBLAS_OP_T;
  CHECK_CUBLAS_ERR(cublasSgemm(cublas_handle_, cuTransB, cuTransA,
      N, M, K, &alpha, B, ldb, A, lda, &beta, C, N));
}

template <>
void caffe_gpu_gemm<double>(const cublasOperation_t TransA,
    const cublasOperation_t TransB, const int M, const int N, const int K,
    const double alpha, const double* A, const double* B, const double beta,
    double* C, cublasHandle_t& cublas_handle_) {
  // Note that cublas follows fortran order.
  int lda = (TransA == CUBLAS_OP_N) ? K : M;
  int ldb = (TransB == CUBLAS_OP_N) ? N : K;
  cublasOperation_t cuTransA =
      (TransA == CUBLAS_OP_N) ? CUBLAS_OP_N : CUBLAS_OP_T;
  cublasOperation_t cuTransB =
      (TransB == CUBLAS_OP_N) ? CUBLAS_OP_N : CUBLAS_OP_T;
  CHECK_CUBLAS_ERR(cublasDgemm(cublas_handle_, cuTransB, cuTransA,
      N, M, K, &alpha, B, ldb, A, lda, &beta, C, N));
}

}} // end of namespace runtime/cuda
